import torch.utils.data as data
import cv2
import os
import random
import image_utils


#  0:neutral, 1:happiness, 2:surprise, 3:sadness, 4:anger, 5:disgust, 6:fear, 7:contempt, 8:unknown, 9:NF

class SFEWDataSet(data.Dataset):
    def __init__(self, SFEW_path, phase, transform=None, basic_aug=False):
        self.phase = phase
        self.transform = transform
        self.SFEW_path = SFEW_path
        type_map = {
            "Angry": 4,  # 197
            "Disgust": 5,  # 53
            "Fear": 6,  # 79
            "Happy": 1,  # 185
            "Neutral": 0,  # 145 15%
            "Sad": 3,  # 162
            "Surprise": 2  # 95
        }

        self.file_paths = []
        self.label = []
        for key in type_map:
            temp_path = os.path.join(self.SFEW_path, key)
            for img in os.listdir(temp_path):
                self.file_paths.append(os.path.join(temp_path, img))
                self.label.append(type_map[key])

        self.basic_aug = basic_aug
        self.aug_func = [image_utils.flip_image, image_utils.add_gaussian_noise]

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        path = self.file_paths[idx]
        image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)  # 读取灰度图
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)  # Gray to RGB
        label = self.label[idx]
        # augmentation
        if self.phase == 'train':
            if self.basic_aug and random.uniform(0, 1) > 0.5:
                index = random.randint(0, 1)
                image = self.aug_func[index](image)

        if self.transform is not None:
            image = self.transform(image)

        return image, label, idx
